package com.idp.saml;

import org.opensaml.Configuration;
import org.opensaml.util.URLBuilder;
import org.opensaml.xml.io.MarshallingException;
import org.opensaml.xml.security.credential.Credential;
import org.opensaml.xml.signature.SignableXMLObject;
import org.opensaml.xml.signature.SignatureConstants;
import org.opensaml.xml.signature.SignatureException;
import org.opensaml.xml.signature.Signer;
import org.opensaml.xml.util.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;

import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.util.List;


public class SignatureMethod {
    private static final Logger logger = LoggerFactory.getLogger(SignatureMethod.class);

    public static byte[] concat(String samlRequest, String relayState, String sigAlg) {
        URLBuilder urlBuilder = new URLBuilder();
        List<Pair<String, String>> queryParams = urlBuilder.getQueryParams();
        queryParams.clear();
        queryParams.add(new Pair<>(AuthnRequestParameterConst.SAML_REQUEST, samlRequest));
        if (relayState != null) {
            queryParams.add(new Pair<>(AuthnRequestParameterConst.RELAY_STATE, relayState));
        }
        queryParams.add(new Pair<>(AuthnRequestParameterConst.SIG_ALG, sigAlg));
        byte[] bytes;
        bytes = urlBuilder.buildQueryString().getBytes(StandardCharsets.UTF_8);
        return bytes;
    }

    public static String extractSignatureAlgorithm(String sigAlg) throws Exception {
        String[] parts = sigAlg.split("#");
        if (parts.length < 2) {
            throw new Exception("unsupported signAlgUrl: " + sigAlg);
        }
        if ("rsa-sha1".equals(parts[1])) {
            return "SHA1withRSA";
        }
        if ("rsa-sha256".equals(parts[1])) {
            return "SHA256withRSA";
        } else {
            throw new Exception("unsupported signAlgUrl: " + sigAlg);
        }
    }

    public static boolean verifyStringSignature(byte[] decodeSignature, byte[] signedData, String signatureAlgorithm, PublicKey publicKey) {
        boolean verify = false;
        try {
            java.security.Signature signature = java.security.Signature.getInstance(signatureAlgorithm);
            signature.initVerify(publicKey);
            signature.update(signedData);
            verify = signature.verify(decodeSignature);
        } catch (java.security.SignatureException | InvalidKeyException | NoSuchAlgorithmException e) {
            logger.error("", e);
            return false;
        }
        return verify;
    }

    public static void signObject(SignableXMLObject signableXMLObject, Credential credential) {
        org.opensaml.xml.signature.Signature signature = SAMLResponse.buildSAMLObject(org.opensaml.xml.signature.Signature.class);
        signature.setSigningCredential(credential);
        signature.setSignatureAlgorithm(Configuration.getGlobalSecurityConfiguration().getSignatureAlgorithmURI(credential));
        signature.setCanonicalizationAlgorithm(SignatureConstants.ALGO_ID_C14N_EXCL_OMIT_COMMENTS);
        try {
            signableXMLObject.setSignature(signature);
            Configuration.getMarshallerFactory().getMarshaller(signableXMLObject).marshall(signableXMLObject);
            Signer.signObject(signature);
        } catch (SignatureException | MarshallingException e) {
            logger.error("sign error", e);
        }
    }
}
